# ==================== 导入库 ====================
import argparse
import warnings
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import roc_auc_score
from imblearn.over_sampling import SMOTE
import optuna
from optuna.samplers import TPESampler
from lightgbm import LGBMClassifier
from lightgbm.callback import early_stopping
import matplotlib.pyplot as plt

# 从公共模块导入辅助函数
from common_group import setup_experiment_paths, save_object, load_data, save_results, save_function_snapshot

# ==================== 配置 ====================
warnings.filterwarnings('ignore')
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
optuna.logging.set_verbosity(optuna.logging.WARNING)


# =====================================================================================
# START OF USER-CONFIGURABLE AREA
# 用户需在此区域定义自己的特征工程、模型和超参数
# =====================================================================================

def define_model():
    """
    在这里定义你的模型.
    
    返回:
        一个未训练的分类器实例 (例如 LGBMClassifier).
    """
    # 示例: 返回一个 LightGBM 分类器
    # 你可以替换成任何你想要的模型, 例如:
    # from sklearn.ensemble import RandomForestClassifier
    # return RandomForestClassifier(random_state=42)
    
    # from xgboost import XGBClassifier
    # return XGBClassifier(random_state=42, use_label_encoder=False, eval_metric='auc')
    
    return LGBMClassifier(random_state=42, objective='binary', metric='auc', n_jobs=-1, verbose=-1)

def define_hyperparameters(trial: optuna.trial.Trial):
    """
    在这里为 Optuna 定义超参数搜索空间.
    
    参数:
        trial (optuna.trial.Trial): Optuna trial 对象.

    返回:
        dict: 包含模型超参数的字典.
    """
    # 示例: LightGBM 的超参数空间
    params = {
        'n_estimators': trial.suggest_int('n_estimators', 200, 2000, step=100),
        'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.2, log=True),
        'num_leaves': trial.suggest_int('num_leaves', 20, 150),
        'max_depth': trial.suggest_int('max_depth', 4, 10),
        'subsample': trial.suggest_float('subsample', 0.6, 1.0),
        'colsample_bytree': trial.suggest_float('colsample_bytree', 0.6, 1.0),
        'reg_alpha': trial.suggest_float('reg_alpha', 1e-8, 10.0, log=True),
        'reg_lambda': trial.suggest_float('reg_lambda', 1e-8, 10.0, log=True),
    }
    return params

def feature_engineering(data: pd.DataFrame, is_train=True, stats=None):
    """
    特征工程 - 此函数逻辑完全基于 common.py 中的 enhanced_feature_engineering.
    """
    print("--- 开始执行自定义特征工程 ---")
    data = data.copy()
    
    # 1. 基础特征清理
    cols_to_drop = ['标准工作时间', '员工编号', '是否超过18岁','性别']
    data.drop(columns=[col for col in cols_to_drop if col in data.columns], inplace=True, errors='ignore')
    
    # 2. 目标编码 & 聚合特征
    if is_train:
        stats = {'月收入_中位数': data['月收入'].median()}
        
        dept_avg_income = data.groupby('所属部门')['月收入'].mean().to_dict()
        role_avg_promo = data.groupby('工作角色')['距离上次升职时长'].mean().to_dict()
        stats['dept_avg_income'] = dept_avg_income
        stats['role_avg_promo'] = role_avg_promo

        data['部门内薪酬水平'] = (data['月收入'] / data['所属部门'].map(dept_avg_income)).fillna(1)
        data['角色内晋升速度'] = (data['距离上次升职时长'] / data['工作角色'].map(role_avg_promo)).fillna(1)
        
        target_col = '是否离职'
        category_features = ['所属部门', '所学专业', '工作角色', '出差频率', '婚姻状况', '是否加班']
        
        for col in category_features:
            if col == '是否加班':
                data[col] = data[col].map({'Yes': 1, 'No': 0})
            mean_dict = data.groupby(col)[target_col].mean().to_dict()
            stats[f'{col}_mean_dict'] = mean_dict
            data[col] = data[col].map(mean_dict)
    else:
        if stats is None: raise ValueError("测试集必须传入从训练集计算的stats字典!")
        
        data['部门内薪酬水平'] = (data['月收入'] / data['所属部门'].map(stats['dept_avg_income'])).fillna(1)
        data['角色内晋升速度'] = (data['距离上次升职时长'] / data['工作角色'].map(stats['role_avg_promo'])).fillna(1)
        
        category_features = ['所属部门', '所学专业', '工作角色', '出差频率', '婚姻状况', '是否加班']
        for col in category_features:
            if col == '是否加班':
                data[col] = data[col].map({'Yes': 1, 'No': 0})

            mean_dict = stats.get(f'{col}_mean_dict', {})
            global_mean = pd.Series(list(mean_dict.values())).mean() if mean_dict else 0.5
            data[col] = data[col].map(mean_dict).fillna(global_mean)
    
    # 3. 核心特征组
    data['低收入标记'] = (data['月收入'] < stats['月收入_中位数']).astype(int)
    data['低满意度'] = (((data['对工作环境满意度'] + data['工作满意度'] + data['与同事的关系'] + data['工作投入度']) / 4) < 2.5).astype(int)
    data['高压力标记'] = ((data['是否加班'] > 0.5) & (data['工作与生活平衡程度'] <= 2)).astype(int)
    data['无激励'] = ((data['股票期权级别'] == 0) & (data['上一次工资增长百分比'] <= 15)).astype(int)
    data['离家距离分箱'] = pd.cut(data['离家距离'], bins=[0, 9, 14, 30], labels=[1, 2, 3], right=True).astype(int)

    # 4. 新增交互特征
    data['逆向工作与生活平衡'] = 5 - data['工作与生活平衡程度'] 
    data['职业倦怠指标'] = (data['是否加班'] * data['逆向工作与生活平衡'] * (5 - data['工作满意度']))
    data['工作稳定性指标'] = data['总工龄'] / (data['工作过的公司数量'] + 1)
    data['晋升停滞指标'] = data['距离上次升职时长'] / (data['职位级别'] + 1)
    data['薪酬竞争力指标'] = data['月收入'] / ((data['职位级别'] + 1) * (data['所学专业'] + 1))
    data['加班且低满意度'] = ((data['是否加班'] > 0.5) & (data['低满意度'] == 1)).astype(int)

    # 5. 清理特征
    cols_to_drop_after_fe = [
        '在目前工作职责的工作年数','目前公司工作年数','上一次工资增长百分比','与同事的关系',
        '总工龄','距离上次升职时长','所属部门','婚姻状况','逆向工作与生活平衡'
    ]
    data.drop(columns=[col for col in cols_to_drop_after_fe if col in data.columns], inplace=True, errors='ignore')

    print("--- 自定义特征工程结束 ---")
    
    if is_train:
        return data, stats
    else:
        return data


def preprocess_data(data: pd.DataFrame):
    """
    数据预处理: 重命名列为中文, 并进行基本清理.
    """
    print("--- 开始执行数据预处理 ---")
    data = data.copy()

    new_cols = [
        '是否离职', '年龄', '出差频率', '所属部门', '离家距离', '教育水平', '所学专业',
        '员工编号', '对工作环境满意度', '性别', '工作投入度', '职位级别', '工作角色',
        '工作满意度', '婚姻状况', '月收入', '工作过的公司数量', '是否超过18岁', '是否加班',
        '上一次工资增长百分比', '绩效评分', '与同事的关系', '标准工作时间', '股票期权级别',
        '总工龄', '上一年培训时长', '工作与生活平衡程度', '目前公司工作年数',
        '在目前工作职责的工作年数', '距离上次升职时长', '跟目前的管理者共事年数'
    ]

    data.columns = new_cols

    print("--- 数据预处理结束 ---")
    return data

# =====================================================================================
# END OF USER-CONFIGURABLE AREA
# 下面的代码是标准化的训练流程,通常不需要修改,仅仅当使用逻辑回归时,请解开注释
# =====================================================================================

class Trainer:
    def __init__(self, group_name, experiment_name, train_data_path, n_trials, use_smote, base_dir='experiments'):
        self.group_name = group_name
        self.experiment_name = experiment_name
        self.train_data_path = train_data_path
        self.n_trials = n_trials
        self.use_smote = use_smote

        self.paths = setup_experiment_paths(
            base_dir=base_dir,
            group_name=self.group_name,
            experiment_name=self.experiment_name
        )
        self.X_train = self.y_train = self.X_val = self.y_val = None
        self.best_model = None
        self.fe_info = None

    def _prepare_data(self):
        print("\n--- 1. 数据加载与预处理 ---")
        raw_data = load_data(self.train_data_path)
        if raw_data is None:
            raise FileNotFoundError("无法加载训练数据，程序终止.")
        
        # 步骤1: 预处理整个数据集
        data = preprocess_data(raw_data)
        if data is None:
            raise ValueError("预处理失败, 请检查数据.")

        # 步骤2: (防数据泄露) 先划分数据集
        train_data, val_data = train_test_split(
            data, test_size=0.25, random_state=42 # 遵从您的修改, 移除 stratify
        )
        
        print("\n--- 2. 对训练集进行特征工程 ---")
        # 步骤3: 对训练集进行特征工程并获得统计数据
        # 此时传入完整的 train_data (包含目标列) 以便进行目标编码
        train_data_fe, stats = feature_engineering(train_data, is_train=True)
        
        print("\n--- 3. 对验证集应用特征工程 ---")
        # 步骤4: 使用从训练集得到的统计数据对验证集进行特征工程
        val_data_fe = feature_engineering(val_data, is_train=False, stats=stats)

        # 步骤5: 分离 X 和 y
        self.y_train = train_data_fe['是否离职']
        self.X_train = train_data_fe.drop(columns=['是否离职'])
        
        self.y_val = val_data_fe['是否离职']
        self.X_val = val_data_fe.drop(columns=['是否离职'])

        # 步骤6: 保存产出物
        print("\n--- 4. 保存特征工程产出物 ---")
        save_object(stats, self.paths['stats'])
        self.fe_info = save_function_snapshot(
            feature_engineering, 
            self.paths['base'], 
            filename='feature_engineering_snapshot.py'
        )

        print(f"训练集形状: {self.X_train.shape}, 验证集形状: {self.X_val.shape}")

        # SMOTE 过采样 (应用于处理后的训练数据)
        if self.use_smote:
            print("\n--- 5. 应用 SMOTE 过采样 ---")
            smote = SMOTE(random_state=42)
            self.X_train, self.y_train = smote.fit_resample(self.X_train, self.y_train)
            print(f"SMOTE后训练集形状: {self.X_train.shape}")
        
        # scaler = StandardScaler()
        # self.X_train = scaler.fit_transform(self.X_train)
        # self.X_val = scaler.transform(self.X_val)
        # save_object(scaler, self.paths['scaler'])
    
    def _objective(self, trial):
        """Optuna 优化目标函数."""
        params = define_hyperparameters(trial)
        model = define_model()
        model.set_params(**params)
        
        # 对于LGBM/XGBoost, 使用 early stopping
        if isinstance(model, LGBMClassifier):
            model.fit(self.X_train, self.y_train,
                        eval_set=[(self.X_val, self.y_val)],
                        callbacks=[early_stopping(100, verbose=False)])
        else:
            model.fit(self.X_train, self.y_train)

        y_pred_proba = model.predict_proba(self.X_val)[:, 1]
        auc = roc_auc_score(self.y_val, y_pred_proba)
        return auc

    def run_hyperparameter_tuning(self):
        print("\n--- 3. 开始超参数优化 ---")
        
        study = optuna.create_study(direction='maximize', sampler=TPESampler(seed=42))
        study.optimize(self._objective, n_trials=self.n_trials, show_progress_bar=True)
        
        print(f"\n优化完成! 最佳AUC: {study.best_value:.4f}")
        print("最佳参数:")
        for key, value in study.best_params.items():
            print(f"  - {key}: {value}")
        
        # 使用最佳参数训练最终模型
        self.best_model = define_model()
        self.best_model.set_params(**study.best_params)
        self.best_model.fit(self.X_train, self.y_train)
        
        # 保存结果
        results = {
            'best_auc': study.best_value,
            'best_params': study.best_params,
            'n_trials': self.n_trials,
            'feature_engineering_info': self.fe_info
        }
        save_results(results, self.paths['results'])

    def save_final_model(self):
        if self.best_model:
            print("\n--- 4. 保存最终模型 ---")
            save_object(self.best_model, self.paths['model'])
        else:
            print("错误: 没有可供保存的模型.")

    def run(self):
        """执行完整的训练流程."""
        self._prepare_data()
        self.run_hyperparameter_tuning()
        self.save_final_model()
        print("\n🎉 训练流程全部完成! 🎉")


def main():
    config = {
        'group_name': 'Local_Debug',    #不用改
        'experiment_name': 'LGBM_Test_Run',#文件前缀
        'train_data_path': '',          # 训练数据路径
        'base_dir': 'experiments',  # 实验结果保存路径
        'n_trials': 10,  # 调试时可适当减小 贝叶斯迭代次数
        'use_smote': True#不用改
    }
    
    print("="*60)
    print("训练配置 (本地调试):")
    for key, value in config.items():
        print(f"  - {key.replace('_', ' ').capitalize()}: {value}")
    print("="*60)

    trainer = Trainer(
        group_name=config['group_name'],
        experiment_name=config['experiment_name'],
        train_data_path=config['train_data_path'],
        n_trials=config['n_trials'],
        use_smote=config['use_smote'],
        base_dir=config['base_dir']
    )
    trainer.run()


if __name__ == '__main__':
    # 使用方法:
    # 直接运行此脚本即可开始训练.
    # 所有配置请在上方 main 函数的 `config` 字典中修改.
    main()
